package com.cxk.fe

import com.cxk.ml.TD_IDF
import org.apache.spark.sql.types.{DoubleType, StringType, StructField, StructType}
import org.apache.spark.sql.{DataFrame, Row, SparkSession, functions}

object LightGBM_FE {

  def lightGBMfe(spark: SparkSession, dir: String, hasTarget: Boolean = false): DataFrame = {
    val reader = spark.read.option("header", "true")
    val (path, target) = if (hasTarget) ("/train.csv", "target") else ("/test.csv", null)
    val df = fe(spark, reader.csv(dir + path), target, song_extra_info(spark, reader.csv(dir + "/song_extra_info.csv")),
      members(spark, reader.csv(dir + "/members.csv"), reader.csv(dir + "/members_init_expiration.csv")),
      reader.csv(dir + "/songs.csv"), reader.csv(dir + "/mean_code_artist_name.csv"), reader.csv(dir + "/artist_index.csv"),
      reader.csv(dir + "/mean_code_composer.csv"), reader.csv(dir + "/composer_index.csv"), reader.csv(dir + "/mean_code_lyricist.csv"),
      reader.csv(dir + "/lyricist_index.csv"), reader.csv(dir + "/mean_code_msno.csv"), reader.csv(dir + "/user_index.csv"),
      reader.csv(dir + "/mean_code_song_id.csv"), reader.csv(dir + "/item_index.csv")
    ).na.fill(-1)

    df.repartition(1).write.option("header", "true").csv(dir + "/lightGBM_fe_" + hasTarget)
    df
  }

  def fe(spark: SparkSession, df: DataFrame, target: String, song_extra_info: DataFrame, members: DataFrame, songs: DataFrame,
         mean_code_artist: DataFrame, artist_index: DataFrame, mean_code_composer: DataFrame, composer_index: DataFrame
         , mean_code_lyricist: DataFrame, lyricist_index: DataFrame, mean_code_msno: DataFrame, user_index: DataFrame
         , mean_code_song_id: DataFrame, item_index: DataFrame): DataFrame = {
    val train_test_df = _0_(df, target).join(song_extra_info, Seq("song_id"), "left").join(members, Seq("msno"), "left")
    val songs_df = song_length(songs)

    val explode = (col: String, d: DataFrame) => d.withColumn(col, functions.explode(functions.split(d(col), "[\\||&|/]")))

    val explode_genre_ids = explode("genre_ids", train_test_df.join(songs_df.select("song_id", "song_length", "genre_ids", "language"),
      Seq("song_id"), "left"))

    val explode_artist_name = explode("artist_name", explode_genre_ids.join(songs_df.select("song_id", "artist_name"), Seq("song_id"), "left"))
    val artist_name = explode_artist_name.join(mean_code_artist, Seq("artist_name"), "left")
      .join(artist_index, Seq("artist_name"), "left").drop("artist_name")

    val explode_composer = explode("composer", artist_name.join(songs_df.select("song_id", "composer"), Seq("song_id"), "left"))
    val composer = explode_composer.join(mean_code_composer, Seq("composer"), "left")
      .join(composer_index, Seq("composer"), "left").drop("composer")

    val explode_lyricist = explode("lyricist", composer.join(songs_df.select("song_id", "lyricist"), Seq("song_id"), "left"))
    val lyricist = explode_lyricist.join(mean_code_lyricist, Seq("lyricist"), "left")
      .join(lyricist_index, Seq("lyricist"), "left").drop("lyricist")

    lyricist.join(mean_code_msno, Seq("msno"), "left").join(user_index, Seq("msno"), "left").drop("msno")
      .join(mean_code_song_id, Seq("song_id"), "left").join(item_index, Seq("song_id"), "left").drop("song_id")
  }

  def song_extra_info(spark: SparkSession, song_extra_info: DataFrame): DataFrame = {
    val rdd = TD_IDF(spark, "song_id", "name").transform(song_extra_info.select("song_id", "name"))
      .select("song_id", "result").rdd.map(row => {
      val vector = row.get(1).toString
      val toDouble = (str: String) => {
        try {
          str.toDouble
        } catch {
          case _: Throwable => System.err.println("Throwable======================" + str); -1d
        }
      }
      val data = vector.substring(1, vector.length - 1).split(",")
      Row(row.get(0), toDouble(data(0)), toDouble(data(1)), toDouble(data(2)))
    })

    spark.createDataFrame(rdd, StructType(List(StructField("song_id", StringType),
      StructField("song_name_vec1", DoubleType), StructField("song_name_vec2", DoubleType), StructField("song_name_vec3", DoubleType))))
  }

  def members(spark: SparkSession, members: DataFrame, members_init_expiration: DataFrame): DataFrame = {
    val gender = exprs("gender", "male", "female")
    val data = members.selectExpr("msno", "city", gender, "CAST(registered_via AS INT) AS registered_via", "CAST(bd AS INT) AS bd")

    val cityCount = data.select("msno", "city").groupBy("city")
      .agg(functions.countDistinct("msno").alias("city_rank"))
    val city_df = FeatureEngineering.zipWithIndex(spark, "city_id", cityCount.orderBy(cityCount("city_rank").desc))

    val bd = "(CASE WHEN bd>0 AND bd<100 THEN bd ELSE -1 END) AS bd"
    data.join(city_df, Seq("city")).drop("city").selectExpr("city_rank", "msno", "city_id", bd, "gender", "registered_via")
      .join(members_init_expiration, Seq("msno"), "left")
  }

  def _0_(df: DataFrame, target: String = null): DataFrame = {
    val source_system_tab = exprs("source_system_tab", "my library", "discover", "search", "radio", "listen with", "explore", "notification", "settings")
    val source_screen_name = exprs("source_screen_name", "Online playlist more", "Radio",
      "Album more", "Search", "Artist more", "Discover Feature", "Discover Chart", "Others profile more",
      "Discover Genre", "My library", "Explore", "Unknown", "Discover New", "Search Trends", "Search Home"
      , "My library_Search", "Self profile more", "Concert", "Payment")
    val source_type = exprs("source_type", "local-library", "online-playlist",
      "local-playlist", "radio", "album", "top-hits-for-artist", "song", "song-based-playlist",
      "listen-with", "topic-article-playlist", "artist", "my-daily-playlist")

    if (target == null) df.selectExpr(source_system_tab, source_screen_name, source_type, "msno", "song_id")
    else df.selectExpr(source_system_tab, source_screen_name, source_type, "msno", "song_id", target)
  }

  def song_length(songs: DataFrame): DataFrame = {
    val _exprs = Array("CAST(song_length AS LONG) AS song_length", "song_id", "genre_ids", "artist_name", "composer", "lyricist", "language")
    val df = songs.selectExpr(_exprs: _*)
    _exprs(0) = "(CASE WHEN song_length>0 AND song_length<60000 THEN 1 " +
      "WHEN song_length>=60000 AND song_length<120000 THEN 2 " +
      "WHEN song_length>=120000 AND song_length<180000 THEN 3 " +
      "WHEN song_length>=180000 AND song_length<240000 THEN 4 " +
      "WHEN song_length>=240000 AND song_length<300000 THEN 5 " +
      "WHEN song_length>=300000 AND song_length<360000 THEN 6 " +
      "WHEN song_length>=360000 AND song_length<420000 THEN 7 " +
      "WHEN song_length>=420000 AND song_length<480000 THEN 8 " +
      "WHEN song_length>=480000 AND song_length<540000 THEN 9 " +
      "WHEN song_length>=540000 AND song_length<600000 THEN 10 " +
      "WHEN song_length>=600000 AND song_length<1200000 THEN 11 " +
      "WHEN song_length>=1200000 AND song_length<1800000 THEN 0 ELSE -1 END) AS song_length"
    df.selectExpr(_exprs: _*)
  }

  def _0_source_system_tab(df: DataFrame, target: String = null): DataFrame = {
    val _exprs = exprs("source_system_tab", "my library", "discover", "search", "radio", "listen with", "explore", "notification", "settings")
    if (target == null) df.selectExpr(_exprs, "msno", "song_id", "source_screen_name", "source_type")
    else df.selectExpr(_exprs, "msno", "song_id", "source_screen_name", "source_type", "target")
  }

  def _0_source_screen_name(df: DataFrame, target: String = null): DataFrame = {
    val _exprs = exprs("source_screen_name", "Online playlist more", "Radio",
      "Album more", "Search", "Artist more", "Discover Feature", "Discover Chart", "Others profile more",
      "Discover Genre", "My library", "Explore", "Unknown", "Discover New", "Search Trends", "Search Home"
      , "My library_Search", "Self profile more", "Concert", "Payment")

    if (target == null) df.selectExpr(_exprs, "msno", "song_id", "source_system_tab", "source_type")
    else df.selectExpr(_exprs, "msno", "song_id", "source_system_tab", "source_type", "target")
  }

  def _0_source_type(df: DataFrame, target: String = null): DataFrame = {
    val _exprs = exprs("source_type", "local-library", "online-playlist",
      "local-playlist", "radio", "album", "top-hits-for-artist", "song", "song-based-playlist",
      "listen-with", "topic-article-playlist", "artist", "my-daily-playlist")

    if (target == null) df.selectExpr(_exprs, "msno", "song_id", "source_system_tab", "source_screen_name")
    else df.selectExpr(_exprs, "msno", "song_id", "source_system_tab", "source_screen_name", "target")
  }

  private[this] def exprs(col: String, values: String*): String = {
    val sb = new StringBuilder()
    for (i <- values.indices) {
      if (i == 0) {
        sb.append("(CASE WHEN ").append(col).append("='").append(values(i)).append("' THEN ").append(i + 1)
      } else {
        sb.append(" WHEN ").append(col).append("='").append(values(i)).append("' THEN ").append(i + 1)
        if (i == values.length - 1)
          sb.append(" ELSE -1 END) AS ").append(col)
      }
    }
    sb.toString()
  }
}
